//#define UPDATE_HOSTS_FROM_SCHEMA_VERSIONS
#define UPDATE_HOSTS_FROM_DESCRIBE_RING
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using Apache.Cassandra;
using Thrift.Protocol;
using Thrift.Transport;

namespace Cassandraemon.Connection
{
	public class ConnectionPool : IDisposable
	{
		readonly object _BorrowReturnLock = new object();

        readonly CassandraConnectionConfig _ConnectionConfig;

        public CassandraConnectionConfig ConnectionConfig
        {
            get { return _ConnectionConfig; }
        }

		public string FirstHost { get { return this._ConnectionConfig.Hosts.First(); } }
		public int FirstPort { get { return this._ConnectionConfig.Port; } }
		public string Keyspace { get { return this._ConnectionConfig.Keyspace; } }
        public bool IsFramed { get { return this._ConnectionConfig.IsFramed; } }
        public TimeSpan Timeout { get { return this._ConnectionConfig.Timeout; } }
        public int TimeoutMillisecond { get { return (int)this._ConnectionConfig.Timeout.TotalMilliseconds; } }
        public int MinPoolCountPerHost { get { return this._ConnectionConfig.MinPoolCountPerHost; } }
        public int MaxPoolCountPerHost { get { return this._ConnectionConfig.MaxPoolCountPerHost; } }
		
		private AutoResetEvent autoResetEvent = new AutoResetEvent(false);

		#region Constructors and Destructors

		public ConnectionPool (string firstHost, 
		                       int firstPort, 
		                       string keyspace)
            : this(new CassandraConnectionConfig(
                CassandraConnectionConfig.SplitHosts(firstHost),
                firstPort,
                keyspace))
		{
		}

		public ConnectionPool (CassandraConnectionConfig config)
		{
            this._ConnectionConfig = config;

			autoResetEvent.Set();

			this.UpdateNodes(config.Hosts);
			this._NodeEnumerator = config.CreateNodeEnumerator(this);
			this.TryUpdateNodes();
		}

		~ConnectionPool()
		{
			Dispose(false);
		}

		#endregion

		#region Pool of ConnectionPool

        static readonly Dictionary<CassandraConnectionConfig, ConnectionPool> _PoolsByConfig =
            new Dictionary<CassandraConnectionConfig, ConnectionPool>(
				CassandraConnectionConfig.PoolKeyComparer);

        internal static ConnectionPool GetOrDefault(CassandraConnectionConfig config)
        {
			ConnectionPool pool;
            if (_PoolsByConfig.TryGetValue(config, out pool) &&
                !pool.IsDisposed)
            {
                return pool;
            }
            return null;
        }

		internal static ConnectionPool GetOrCreate(CassandraConnectionConfig config)
		{
            if (config == null)
            {
                throw new ArgumentNullException();
            }

            ConnectionPool pool;
            if (_PoolsByConfig.TryGetValue(config, out pool) &&
                !pool.IsDisposed)
            {
                return pool;
            }

			lock (_PoolsByConfig)
			{
                CleanupDisposedPoolsUnsafe();

                if (_PoolsByConfig.TryGetValue(config, out pool))
                {
                    return pool;
                }

				foreach (var kv in _PoolsByConfig)
				{
					pool = kv.Value;
					if (pool.Nodes.Intersect(config.Hosts).Any() &&
						pool._ConnectionConfig.EqualsPoolKeysExceptHosts(config))
					{
						return pool;
					}
				}

				pool = new ConnectionPool(config);
                _PoolsByConfig.Add(config, pool);
				return pool;
			}
		}

        static void CleanupDisposedPoolsUnsafe()
        {
            var rgcfgDisposed = _PoolsByConfig
                .Where(kv => kv.Value.IsDisposed)
                .Select(kv => kv.Key)
                .ToList();
            foreach (var cfg in rgcfgDisposed)
            {
                _PoolsByConfig.Remove(cfg);
            }
        }

#if NOT_USED
		internal static void ClearPools()
		{
			lock (Pools)
			{
				foreach (var p in Pools.Values)
				{
					p.Close();
				}

				Pools.Clear();
			}
		}
#endif

		#endregion

		#region Clients by host

		TimeSpan _DisableConnectionAfterError = TimeSpan.FromMinutes(5);

		public TimeSpan DisableConnectionAfterError
		{
			get { return _DisableConnectionAfterError; }
			set { _DisableConnectionAfterError = value; }
		}

		[DebuggerDisplay("{Host}({IdleClients.Count}/{Clients.Count})")]
		class HostClients
		{
            readonly ConnectionPool _ConnectionPool;
            public ConnectionPool ConnectionPool { get { return _ConnectionPool; } } 
			public readonly string Host;
			public int Port { get { return this._ConnectionPool.FirstPort; } }
			public readonly List<Cassandra.Client> Clients;
			public readonly Queue<Cassandra.Client> IdleClients;

			public HostClients(
				ConnectionPool pool,
				string host)
			{
				this._ConnectionPool = pool;
				this.Host = host;
				var maxConnections = pool.MaxPoolCountPerHost;
				var capacity = Math.Min(32, maxConnections);
				this.Clients = new List<Cassandra.Client>(capacity);
				this.IdleClients = new Queue<Cassandra.Client>(capacity);
			}

			public int Count
			{
				get { return this.Clients.Count; }
			}

			public override string ToString()
			{
				return string.Format("{0}({1}/{2})",
					this.Host,
					this.IdleClients.Count,
					this.Clients.Count);
			}

			public Cassandra.Client Add()
			{
                var sw = Stopwatch.StartNew();
                var config = this._ConnectionPool.ConnectionConfig;
                var host = this.Host;
                var port = config.Port;
				TTransport socket = new TSocket(host, port);
			
				if(config.IsFramed)
				{
                    socket = new TFramedTransport(socket);
				}
				var protocol = new TBinaryProtocol(socket);
				var client = new Cassandra.Client(protocol);
				try
				{
					socket.Open();
					client.set_keyspace(config.Keyspace);
				}
				catch (Exception err)
				{
                    OnConnectFailed(this, sw, err);
					this.MarkError(err);
					client.Close();
					throw;
				}
                OnConnected(this, sw);

				this.Clients.Add(client);
				return client;
			}

			public Cassandra.Client Borrow(
                int maxConnections,
                bool fIgnoreRecentError)
			{
				// Find an idle connection.

				Cassandra.Client client;
				var queue = this.IdleClients;
				if (queue.Count > 0)
				{
					client = queue.Dequeue();
					try
					{
						var transport = client.InputProtocol.Transport;
						if (!transport.IsOpen)
						{
							transport.Open();
						}
					}
					catch (Exception err)
					{
						CassandraContext.TraceError("Borrow.Dequeue: {0}: {1}",
							this.Host,
							err);
						this.MarkError(err);
						client.Close();
						throw;
					}
					return client;
				}

				if (!fIgnoreRecentError && this.IsLastErrorRecent)
				{
					return null;
				}

				if (maxConnections > 0 && this.Clients.Count >= maxConnections)
				{
					return null;
				}

				client = this.Add();
				return client;
			}

			#region Errors

			private Exception _LastError;
			private DateTime _LastErrorAtUtc;

			public bool IsLastErrorRecent
			{
				get
				{
					var lastErrorAtUtc = this._LastErrorAtUtc;
					if (lastErrorAtUtc > DateTime.MinValue)
					{
						var sinceLastError = DateTime.UtcNow.Subtract(lastErrorAtUtc);
						if (sinceLastError < this._ConnectionPool._DisableConnectionAfterError)
						{
							return true;
						}
						this._LastErrorAtUtc = DateTime.MinValue;
					}
					return false;
				}
			}

			public void MarkError(Exception err)
			{
				this._LastError = err;
				this._LastErrorAtUtc = DateTime.UtcNow;
			}

			#endregion

			#region Close

			public void CloseAllIdleClients()
			{
				var queue = this.IdleClients;
				foreach (var client in queue)
				{
					client.Close();
				}
				queue.Clear();
			}

			public void CloseAllClients()
			{
				var list = this.Clients;
				foreach (var client in list)
				{
					client.Close();
				}
				list.Clear();
			}

			#endregion
		}

		private Dictionary<string, HostClients> _ClientsByHost
			= new Dictionary<string, HostClients>();

		public ICollection<string> Nodes
		{
			get { return this._ClientsByHost.Keys; }
		}

		[Obsolete("Use Nodes property because pluggable NodeEnumerator makes it not a TokenMap anymore")]
		public ICollection<string> TokenMap
		{
			get { return this.Nodes; }
		}

        HostClients HostClientsFromHostName(string host)
        {
            HostClients rgcli;
            if (!this._ClientsByHost.TryGetValue(host, out rgcli))
            {
                CassandraContext.TraceInformation("Host \"{0}\" added", host);
                rgcli = new HostClients(this, host);
                this._ClientsByHost.Add(host, rgcli);
            }
            return rgcli;
        }

        void UpdateHostClients(IEnumerable<string> hosts)
        {
			// Add new token to pool

			var rghostMissing = new HashSet<string>(this._ClientsByHost.Keys);
			foreach (var host in hosts)
			{
				if (!this._ClientsByHost.ContainsKey(host))
				{
					CassandraContext.TraceInformation("Host \"{0}\" added", host);
                    var rgcli = new HostClients(this, host);
                    this._ClientsByHost.Add(host, rgcli);
                }
				rghostMissing.Remove(host);
			}

            // delete missing token from pool

            foreach (var host in rghostMissing)
            {
                CassandraContext.TraceInformation("Host \"{0}\" removed", host);

                HostClients rgcli;
                if (this._ClientsByHost.TryGetValue(host, out rgcli))
                {
                    // Close client stored in idle pool. Lent client will be closed at return.

                    rgcli.CloseAllIdleClients();

                    this._ClientsByHost.Remove(host);
                }
            }
        }

		#endregion

		#region Nodes

		readonly ICassandraNodeEnumerator _NodeEnumerator;

		public ICassandraNodeEnumerator NodeEnumerator
		{
			get { return _NodeEnumerator; }
		}

		private void TryUpdateNodes()
		{
			try
			{
				this._NodeEnumerator.UpdateNodes();
			}
			catch (Exception err)
			{
				CassandraContext.TraceError("UpdateNodes {0}",
					err);
			}
		}

		public void UpdateNodes(IEnumerable<string> nodes)
		{
			lock (this._BorrowReturnLock)
			{
				this.UpdateHostClients(nodes);
			}
		}

		#endregion

		#region Borrow/Return

		public Cassandra.Client Borrow()
		{
			if(disposed) throw new ObjectDisposedException("ConnectionPool is already disposed.");

			var start = DateTime.UtcNow;
			for (; ; )
			{
				lock (this._BorrowReturnLock)
				{
					// Try to get an idle connection,
                    // or new connection up to MinPoolCountPerHost.

					var rghost = this._NodeEnumerator.NextNodes();
                    var client = this.BorrowCore(rghost, this.MinPoolCountPerHost);
					if (client != null)
					{
						return client;
					}

					// If fail to get Connection, add new it when pool have more capacity.

					var rgcli = this._ClientsByHost.Values
						.OrderBy(cli => cli.Count)
						.ToList();
					client = this.BorrowCore(rgcli, this.MaxPoolCountPerHost);
					if (client != null)
					{
						return client;
					}
				}

				var elapsed = DateTime.UtcNow.Subtract(start);
				var timeUntilTimeout = this.Timeout.Subtract(elapsed);
                if (timeUntilTimeout <= TimeSpan.Zero)
                {
                    break;
                }

				CassandraContext.TraceWarning(
					"Waiting for {0} until connections are returned (max={1})",
					timeUntilTimeout,
					this.MaxPoolCountPerHost);
				autoResetEvent.WaitOne(timeUntilTimeout);
			}

			var rgstatus = this._ClientsByHost
				.Select(x => x.Value.ToString())
				.ToArray();
			throw new TimeoutException(string.Format(
				"Timeout! Cassandra connections are used to other thread more than {0}: {1}",
				this.Timeout,
				string.Join(", ", rgstatus)));
		}

        Cassandra.Client BorrowCore(
            IEnumerable<string> rghost,
            int maxConnections)
        {
            var rgrgcli = rghost
                .Select<string, HostClients>(this.HostClientsFromHostName)
                .ToList();
            return this.BorrowCore(rgrgcli, maxConnections);
        }

		Cassandra.Client BorrowCore(
			ICollection<HostClients> rgrgcli,
			int maxConnections)
		{
			var fAllRecentError = rgrgcli.All(rgcli => rgcli.IsLastErrorRecent);

			var cError = 0;
			foreach (var rgcli in rgrgcli)
			{
				try
				{
					var client = rgcli.Borrow(maxConnections, fAllRecentError);
					if (client != null)
					{
						return client;
					}
				}
				catch (Exception)
				{
					// Re-throw if all hosts failed,
					// which probably means waiting until timeout will not resolve the problem.

					cError++;
					if (cError == rgrgcli.Count)
					{
						throw;
					}
				}
			}

			return null;
        }

        #endregion

        #region Return/Errors

		public void Return(Cassandra.Client client)
		{
			if(disposed) throw new ObjectDisposedException("ConnectionPool is already disposed.");
			
			TSocket socket = client.OutputProtocol.Transport.GetTSocket();
			
			if(socket == null) 
				throw new ArgumentException("Error, could not get client socket.");
				
			var poolKey = socket.Host;
				
			// if socket is closed and can't open, then Update token.
			if(!socket.IsOpen)
			{
				try
				{
					socket.Open();
				}
				catch
				{
					socket.Close();
					
					TryUpdateNodes();
					
					return;
				}
			}
			
			
//			// If there are enough connection, discard returned connection.
//			lock(idlePool)
//			lock(allPool)
//			{
//				try
//				{
//					if(idlePool.Any(x => x.Value.Count != 0))
//					{
//						if(allPool[poolKey].Count > MinPoolCountPerHost)
//						{
//							try
//							{
//								socket.Close();
//							}
//							finally
//							{
//								allPool[poolKey].Remove(client);
//							}
//							
//							return;
//						}
//					}
//				}
//				catch
//				{
//					// If not catch exception, then Lock may be not released.
//				}
//			}
			
			// If host exists, then Enqueue, else socket close.

			lock (this._BorrowReturnLock)
			{
				HostClients rgcli;
				if (!this._ClientsByHost.TryGetValue(poolKey, out rgcli))
				{
					client.Close();
					return;
				}

				rgcli.IdleClients.Enqueue(client);
				autoResetEvent.Set();
			}
		}
		
		public void ErrorReturn(Cassandra.Client client, Exception err)
		{
			if(disposed) throw new ObjectDisposedException("ConnectionPool is already disposed.");
			
			TSocket socket = client.OutputProtocol.Transport.GetTSocket();
			
			if(socket == null) 
				throw new ArgumentException("Error, could not get client socket.");
				
			var poolKey = socket.Host;

			client.Close();
			
			// Update token

			HostClients rgcli;
			if (this._ClientsByHost.TryGetValue(poolKey, out rgcli))
			{
				rgcli.MarkError(err);
				TryUpdateNodes();
			}
		}

		#endregion

		#region Close/Dispose

		private bool disposed = false;

		internal bool IsDisposed { get { return disposed; } }

		private void CloseAllClient()
		{
			var rgcliByHost = this._ClientsByHost;
			if(rgcliByHost != null)
			{
				foreach(var rgcli in rgcliByHost.Values)
				{
					rgcli.CloseAllClients();
				}
			}
		}
		
		public void Close()
		{
			Dispose();
		}
		
		public void Dispose ()
		{
			Dispose(true);
			
			GC.SuppressFinalize(this);
		}
		
		private void Dispose(bool disposing)
		{
			if(!this.disposed)
			{
				if(disposing)
				{
					if(autoResetEvent != null) autoResetEvent.Close();
					CloseAllClient();
				}
				
				autoResetEvent = null;
				this._ClientsByHost = null;

				disposed = true;
			}
		}

		#endregion

        #region Events/Trace

        static void OnConnected(HostClients hc, Stopwatch sw)
        {
            var elapsed = sw.Elapsed;
            CassandraContext.TraceInformation("Connected: {0}: {1}",
                hc.Host,
                elapsed);

            var handler = Connected;
            if (handler != null)
            {

                var e = new CassandraConnectionEventArgs(hc.Host, elapsed, null);
                handler(hc.ConnectionPool, e);
            }
        }

        static void OnConnectFailed(HostClients hc, Stopwatch sw, Exception err)
        {
            CassandraContext.TraceError("Connection failed: {0}: {1}",
                hc.Host,
                err);

            var handler = ConnectionFailed;
            if (handler != null)
            {
                var elapsed = sw.Elapsed;
                var e = new CassandraConnectionEventArgs(hc.Host, elapsed, err);
                handler(hc.ConnectionPool, e);
            }
        }

        public static event EventHandler<CassandraConnectionEventArgs> Connected;
        public static event EventHandler<CassandraConnectionEventArgs> ConnectionFailed;

        #endregion
    }
}
